#include "DirectPossionSolver.H"

void DirectPossionSolver<2>::doFindUnknownMap()
{
    Box<2> ghostedBx = emGrid.inflate(1);
    unknownMap.resize(ghostedBx);
    unknownMap = -1;
    numUnknownsAll = 0;
    // Step 1. Find the one-to-one correspondences between the cells and the unknowns.
    loop_box_2(emGrid, i, j)
    {
        if (emGrid.inDomain(i, j))
        {
            unknownMap(i, j) = numUnknownsAll;
            ++numUnknownsAll;
        }
    }
    numUnknownsInDomain = numUnknownsAll;
    // 1.1  handle the regualr ghost cells (only Neumann 's ghost cell is valid)
    for (int d = 0; d < Dim; ++d)
    {
        Box<1> sideBox = reduce(emGrid, d);
        for (int side : {-1, 1})
            if (bcTypes[2 * d + (side + 1) / 2] == 'N')
            {
                int a = (side == -1) ? (emGrid.lo()[d]) : (emGrid.hi()[d]);
                a += side;
                loop_box_1(sideBox, j)
                {
                    iVec idx;
                    idx[d] = a;
                    idx[1 - d] = j;
                    unknownMap(idx) = numUnknownsAll++;
                }
            }
    }
    numUnknownsExceptIrGhost = numUnknownsAll;
    // 1.2 handle the irregular ghost cell
    const Box<2> miniBox = emGrid.getMiniBox();
    loop_box_2(miniBox, i, j)
    {
        if (!emGrid.inDomain(i, j))
        {
            int m[] = {emGrid.inDomain(i, j),
                       emGrid.inDomain(i - 1, j), emGrid.inDomain(i + 1, j),
                       emGrid.inDomain(i, j - 1), emGrid.inDomain(i, j + 1)};
            if (std::any_of(&m[0], &m[5], [](int x)
                            { return x == true; }))
            {
                unknownMap(i, j) = numUnknownsAll;
                ++numUnknownsAll;
            }
        }
    }
    unknownMapInv.resize(numUnknownsAll);
    loop_box_2(ghostedBx, i, j) if (unknownMap(i, j) != -1)
        unknownMapInv[unknownMap(i, j)] = (iVec){i, j};
};

void DirectPossionSolver<2>::doFillMat()
{
    // Step 2. Fill in the coefficients of matrix.
    const auto dx = emGrid.spacing();
    poissonMat.resize(numUnknownsAll);
    poissonMat = 0.0;
    iVec lo = emGrid.lo(), hi = emGrid.hi();
    // 2.1. Fill Matrix for the regular boundary and regular ghost. Also generate standard stencil.
    for (int d = 0; d < Dim; ++d)
    {
        Box<1> sideBox = reduce(emGrid, d);
        for (int side : {-1, 1})
        {
            int a = (side == -1) ? (emGrid.lo()[d]) : (emGrid.hi()[d]);
            switch (bcTypes[2 * d + (side + 1) / 2])
            {
            case 'D': // For 'D', fill the boundary.
                loop_box_1(sideBox, j)
                {
                    iVec idx;
                    idx[d] = a;
                    idx[1 - d] = j;
                    int m = unknownMap(idx);
                    poissonMat(m, m) = 1;
                }
                (side == -1) ? (lo[d] -= side) : (hi[d] -= side); // standard Box中不包括'D'边界
                break;
            case 'N': // For 'N', fill the ghost.
                a += side;
                loop_box_1(sideBox, j)
                {
                    iVec idx;
                    idx[d] = a;
                    idx[1 - d] = j;
                    int m[] = {unknownMap(idx), unknownMap(idx - iVec::unit(d) * 2 * side)};
                    poissonMat(m[0], m[0]) = 1.0 / (2 * dx[1 - d]);
                    poissonMat(m[0], m[1]) = -1.0 / (2 * dx[1 - d]);
                }
                break;
            default:
                break;
            }
        }
    }
    // 2.2  Fill Matrix for standard interior, i.e. can use the standard five point difference scheme.
    Box<2> standardBox(lo, hi);
    loop_box_2(standardBox, i, j)
    {
        if ((unknownMap(i, j) < numUnknownsExceptIrGhost) && (unknownMap(i, j) > -1))
        {
            int m[] = {unknownMap(i, j),
                       unknownMap(i - 1, j), unknownMap(i + 1, j),
                       unknownMap(i, j - 1), unknownMap(i, j + 1)};
            poissonMat(m[0], m[0]) = 2.0 / (dx[0] * dx[0]) + 2.0 / (dx[1] * dx[1]);
            poissonMat(m[0], m[1]) = poissonMat(m[0], m[2]) = (-1.0) / (dx[0] * dx[0]);
            poissonMat(m[0], m[4]) = poissonMat(m[0], m[3]) = (-1.0) / (dx[1] * dx[1]);
        }
    }
    // 2.4 Fill Matrix for Irregualr ghost.
    switch (bcTypes[2 * Dim])
    {
    case 'D':
    {
        for (int i = numUnknownsExceptIrGhost; i < numUnknownsAll; ++i)
            doFillMat_DiriIrBdr(unknownMapInv[i]);
        break;
    }
    case 'N':
    {
        for (int i = numUnknownsExceptIrGhost; i < numUnknownsAll; ++i)
            doFillMat_NeumIrBdr(unknownMapInv[i]);
        break;
    }
    default:
        break;
    }
};

template <class TFunc>
const Tensor<Real, 1> DirectPossionSolver<2>::doFillRhs(const Tensor<Real, Dim - 1> *aRectBdrData, const Tensor<Real, Dim> &aDomainData, const TFunc &aIrBdrDataExpr) const
{
    // Step 2. Fill in the coefficients of matrix and rhs.
    const auto dx = emGrid.spacing();
    Tensor<Real, 1> rhs;
    rhs.resize(numUnknownsAll);
    // 2.3 Fill all rhs except irBdr.
    loop_box_1(rhs.box(), i)
    {
        // 首先分配好 DDphi.
        if (i < numUnknownsInDomain)
            rhs(i) = aDomainData(unknownMapInv[i]);
        // 其次分配规则边界。
        for (int d = 0; d < Dim; ++d)
        {
            Box<1> sideBox = reduce(emGrid, d);
            for (int side : {-1, 1})
            {
                int a = (side == -1) ? (emGrid.lo()[d]) : (emGrid.hi()[d]);
                switch (bcTypes[2 * d + (side + 1) / 2])
                {
                case 'D': // For 'D', fill the boundary.
                    loop_box_1(sideBox, j)
                    {
                        iVec idx;
                        idx[d] = a;
                        idx[1 - d] = j;
                        int m = unknownMap(idx);
                        rhs(m) = aRectBdrData[2 * d + (side + 1) / 2](j);
                    }
                    break;
                case 'N': // For 'N', fill the ghost.
                    a += side;
                    loop_box_1(sideBox, j)
                    {
                        iVec idx;
                        idx[d] = a;
                        idx[1 - d] = j;
                        int m = unknownMap(idx);
                        rhs(m) = aRectBdrData[2 * d + (side + 1) / 2](j);
                    }
                    break;
                default:
                    break;
                }
            }
        }
    }
    // 2.4 Fill RHS for Irregualr ghost.
    switch (bcTypes[2 * Dim])
    {
    case 'D':
        for (int i = numUnknownsExceptIrGhost; i < numUnknownsAll; ++i)
        {
            const iVec idx = unknownMapInv[i];
            bool Filled = false;
            for (int d = 0; d < Dim; ++d)
            {
                for (int side : {-1, 1})
                {
                    if (emGrid.inDomain(idx + iVec::unit(d) * side))
                    {
                        iVec idx_s = idx + iVec::unit(d) * ((side - 1) / 2);
                        rVec cutP = emGrid.cutPointMap(idx_s, d);
                        rhs(unknownMap(idx)) = aIrBdrDataExpr(cutP);
                        Filled = true;
                    }
                    if (Filled)
                        break;
                }
                if (Filled)
                    break;
            }
        }
        break;
    case 'N':
        for (int i = numUnknownsExceptIrGhost; i < numUnknownsAll; ++i)
        {
            const iVec idx = unknownMapInv[i];
            rVec x = idx * dx;
            Real rad = emGrid.embeddingCircle.getRadian(x);
            rVec cutP = emGrid.embeddingCircle.polarRep(rad);
            assert(dx[0] == dx[1]);
            rhs(unknownMap(idx)) = aIrBdrDataExpr(cutP);
        }
        break;
    default:
        break;
    }
    return rhs;
};

const Tensor<Real, 1> DirectPossionSolver<2>::doDirectSolve(const Tensor<Real, Dim - 1> &rhs) const
{
    Tensor<Real, 1> sol;
    Tensor<int, 1> ipiv;
    Tensor<Real, 2> poissonFact;
    sol = rhs.copy();
    poissonFact = poissonMat.copy();
    ipiv.resize(numUnknownsAll);
    auto info = LAPACKE_dgesv(LAPACK_COL_MAJOR, numUnknownsAll /*n*/, 1 /*nrhs*/,
                              poissonFact.data() /*a*/, numUnknownsAll /*lda*/, ipiv.data() /*ipiv*/, sol.data(), numUnknownsAll);
    if (info != 0)
        throw std::runtime_error("[DirectSolver::DirectSolver]fails.");
    return sol;
}

void DirectPossionSolver<2>::doFillMat_DiriIrBdr(const iVec &idx)
{
    for (int d = 0; d < Dim; ++d)
    {
        for (int side : {-1, 1})
        {
            if (emGrid.inDomain(idx + iVec::unit(d) * side))
            {
                iVec idx_s = idx + iVec::unit(d) * ((side - 1) / 2);
                Real cutRadio = emGrid.cutRadios[d](idx_s);
                rVec cutP = emGrid.cutPointMap(idx_s, d);
                switch (ordIntpDiriIrBdr)
                {
                case 1:
                {
                    int m[] = {unknownMap(idx),
                               unknownMap(idx + iVec::unit(d) * side)};
                    if(side == 1)
                    {
                        poissonMat(m[0], m[0]) = 1 - cutRadio;
                        poissonMat(m[0], m[1]) = cutRadio;                    
                    }
                    else
                    {
                        poissonMat(m[0], m[1]) = 1 - cutRadio;
                        poissonMat(m[0], m[0]) = cutRadio;                    
                    }
                    return;
                }
                case 2:
                {
                    int m[] = {unknownMap(idx),
                               unknownMap(idx + iVec::unit(d) * side), unknownMap(idx + iVec::unit(d) * side * 2)};
                    if (side == 1)
                    {
                        poissonMat(m[0], m[0]) = (cutRadio - 1) * (cutRadio - 2) / 2.;
                        poissonMat(m[0], m[1]) = -(cutRadio - 2) * cutRadio;
                        poissonMat(m[0], m[2]) = (cutRadio - 1) * cutRadio / 2.;
                    }
                    else
                    {
                        poissonMat(m[0], m[0]) = (cutRadio + 1) * cutRadio / 2.;
                        poissonMat(m[0], m[1]) = -(cutRadio - 1) * (cutRadio + 1);
                        poissonMat(m[0], m[2]) = (cutRadio - 1) * cutRadio / 2.;
                    }
                    return;
                }
                default:
                    break;
                }
            }
        }
    }
}

void DirectPossionSolver<2>::doFillMat_NeumIrBdr(const iVec &idx)
{
    const auto dx = emGrid.spacing();
    rVec x = idx * dx;
    Real rad = emGrid.embeddingCircle.getRadian(x);
    rVec cutP = emGrid.embeddingCircle.polarRep(rad);
    assert(dx[0] == dx[1]);
    Real l = dx[0] / std::max(abs(sin(rad)), abs(cos(rad)));
    Real cut_radio = std::min(abs(tan(rad)), 1.0 / abs(tan(rad)));
    vector<iVec> allCorners = {{-1, -1}, {1, -1}, {1, 1}, {-1, 1}};
    vector<iVec> allFaceMids = {{1, 0}, {0, 1}, {-1, 0}, {0, -1}};
    int idxCorner, idxFaceMid;
    if ((rad >= -M_PI) && (rad <= -M_PI / 2.0))
        idxCorner = 0;
    else if ((rad >= -M_PI / 2.0) && (rad <= 0.0))
        idxCorner = 1;
    else if ((rad >= 0.0) && (rad <= M_PI / 2.0))
        idxCorner = 2;
    else
        idxCorner = 3;
    if (abs(rad) > 3 * M_PI / 4.0)
        idxFaceMid = 2;
    else if (abs(rad) < M_PI / 4.0)
        idxFaceMid = 0;
    else if (rad > 0)
        idxFaceMid = 1;
    else
        idxFaceMid = 3;
    switch (ordIntpNeumIrBdr)
    {
    case 1:
    {
        int m[] = {unknownMap(idx),
                   unknownMap(idx + allCorners[idxCorner]), unknownMap(idx + allFaceMids[idxFaceMid])};
        poissonMat(m[0], m[0]) = 1.0 / l;
        poissonMat(m[0], m[1]) = -cut_radio / l;
        poissonMat(m[0], m[2]) = -(1.0 - cut_radio) / l;
        return;
    }
    case 2:
    {
        iVec Corner1 = idx + allCorners[idxCorner];
        iVec FaceMid1 = idx + allFaceMids[idxFaceMid];
        iVec addPoint1 = Corner1 * 2 - FaceMid1;
        iVec Corner2 = Corner1 * 2 - idx;
        iVec FaceMid2 = FaceMid1 * 2 - idx;
        iVec addPoint2 = Corner1 + FaceMid1 - idx;
        Real a = norm(x - cutP);
        int n[] = {unknownMap(idx),
                   unknownMap(FaceMid1), unknownMap(Corner1), unknownMap(addPoint1),
                   unknownMap(FaceMid2), unknownMap(addPoint2), unknownMap(Corner2)};
        poissonMat(n[0], n[0]) = (3.0 * l - 2.0 * a) / (2.0 * l * l);
        poissonMat(n[0], n[1]) = -(2.0 * (l - a) / (l * l)) * (0.5 * (cut_radio - 2) * (cut_radio - 1));
        poissonMat(n[0], n[2]) = -(2.0 * (l - a) / (l * l)) * (cut_radio * (2 - cut_radio));
        poissonMat(n[0], n[3]) = -(2.0 * (l - a) / (l * l)) * (0.5 * cut_radio * (cut_radio - 1));
        poissonMat(n[0], n[4]) = -((2.0 * a - l) / (2.0 * l * l)) * (0.5 * (2 * cut_radio - 2) * (2 * cut_radio - 1));
        poissonMat(n[0], n[5]) = -((2.0 * a - l) / (2.0 * l * l)) * (2 * cut_radio * (2 - 2 * cut_radio));
        poissonMat(n[0], n[6]) = -((2.0 * a - l) / (2.0 * l * l)) * (0.5 * 2 * cut_radio * (2 * cut_radio - 1));
        return;
    }
    default:
        break;
    }
}

template <class TFunc>
const Real DirectPossionSolver<2>::computeErrNorm(const Tensor<Real, 1> &sol, const TFunc &exactSolution, const int q) const
{
    const auto dx = emGrid.spacing();
    Tensor<Real, 1> exact, err;
    exact.resize(numUnknownsInDomain), err.resize(numUnknownsInDomain);
    exact = 0, err = 0;
    Real err_norm = 0;
    assert(q >= 0);
    for (int i = 0; i < numUnknownsInDomain; i++)
    {
        iVec idx = unknownMapInv[i];
        exact(i) = exactSolution(idx * dx);
        err(i) = sol(i) - exact(i);
        if (q == 0)
            err_norm = std::max(err_norm, abs(err(i)));
        else
            err_norm += pow(abs(err(i)), q);
    }
    if (q != 0)
        err_norm = pow((dx[0] * dx[1]) * err_norm, 1.0 / q);
    return err_norm;
}

template <class TFunc>
const void DirectPossionSolver<2>::plotIn(const Tensor<Real, 1> &sol, const std::string &_file, const TFunc &exactSolution) const
{
    const auto dx = emGrid.spacing();
    std::ofstream os(_file);
    Tensor<Real, 2> solution(emGrid), exact(emGrid), err(emGrid);
    solution = 0, exact = 0, err = 0;
    for (int i = 0; i < numUnknownsInDomain; i++)
    {
        iVec idx = unknownMapInv[i];
        solution(idx) = sol(i);
        exact(idx) = exactSolution(idx * dx);
        err(idx) = solution(idx) - exact(idx);
    }
    os << "%% dx = " << dx << "\n"
       << "solution = " << solution << ";\n"
       << "exact = " << exact << ";\n"
       << "err = " << err << ";\n";
    os.close();
}
